import numpy as np
import matplotlib.pyplot as plt

class KMeans:
    def __init__(self, k=2, max_iter=100, tol=1e-4):
        self.k = k  # 聚类数量
        self.max_iter = max_iter  # 最大迭代次数
        self.tol = tol  # 质心变化阈值（收敛条件）
        self.centroids = None  # 质心
        self.labels = None  # 样本所属簇标签

    def euclidean_distance(self, x1, x2):
        """计算欧氏距离"""
        return np.sqrt(np.sum((x1 - x2) **2))

    def fit(self, X):
        # 随机初始化质心（从样本中选择k个）
        np.random.seed(42)  # 固定随机种子，保证结果可复现
        self.centroids = X[np.random.choice(X.shape[0], self.k, replace=False)]

        for _ in range(self.max_iter):
            # 步骤1：分配样本到最近的质心
            self.labels = np.zeros(X.shape[0])
            for i, x in enumerate(X):
                distances = [self.euclidean_distance(x, centroid) for centroid in self.centroids]
                self.labels[i] = np.argmin(distances)  # 距离最小的簇索引

            # 步骤2：更新质心（计算每个簇的均值）
            new_centroids = np.zeros((self.k, X.shape[1]))
            for c in range(self.k):
                cluster_points = X[self.labels == c]
                new_centroids[c] = np.mean(cluster_points, axis=0)

            # 检查是否收敛（质心变化小于阈值）
            centroid_shift = np.sum(np.abs(new_centroids - self.centroids))
            if centroid_shift < self.tol:
                break
            self.centroids = new_centroids

    def predict(self, X):
        """预测新样本的簇标签"""
        labels = np.zeros(X.shape[0])
        for i, x in enumerate(X):
            distances = [self.euclidean_distance(x, centroid) for centroid in self.centroids]
            labels[i] = np.argmin(distances)
        return labels

# 示例：生成模拟数据并测试
if __name__ == "__main__":
    # 生成3个簇的模拟数据
    np.random.seed(42)
    cluster1 = np.random.normal(loc=[2, 2], scale=0.5, size=(100, 2))  # 簇1：中心(2,2)
    cluster2 = np.random.normal(loc=[5, 5], scale=0.5, size=(100, 2))  # 簇2：中心(5,5)
    cluster3 = np.random.normal(loc=[8, 2], scale=0.5, size=(100, 2))  # 簇3：中心(8,2)
    X = np.vstack([cluster1, cluster2, cluster3])  # 合并为300个样本的数据集

    # 训练k-means（k=3）
    kmeans = KMeans(k=3)
    kmeans.fit(X)

    # 可视化结果
    plt.scatter(X[:, 0], X[:, 1], c=kmeans.labels, cmap='viridis', alpha=0.6)
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], c='red', marker='X', s=200, label='质心')
    plt.xlabel('特征1')
    plt.ylabel('特征2')
    plt.title('k-means聚类结果（k=3）')
    plt.legend()
    plt.show()